import numpy as np
import os
from scipy.signal import hilbert
# 参数定义
sampling_rate = 128       # DEAP采样率 (Hz)
window_length = 3         # 窗口长度 (秒)
overlap = 2.5             # 重叠时长 (秒)
n_trials = 40             # 试次数
n_channels = 32           # 通道数
n_labels = 4              # 标签维度（例如valence, arousal, dominance, liking）

# 计算窗口参数
window_samples = window_length * sampling_rate    # 384
step_samples = int((window_length - overlap) * sampling_rate)  # 64 (0.5秒)
n_segments_per_trial = (60 * sampling_rate -
                        window_samples) // step_samples + 1  # 115


# --------------------------------------------------
# 步骤2：分割EEG数据
# --------------------------------------------------
def segment_trial(trial_eeg, window_samples, step_samples):
    """
    输入单个试次的EEG数据 (形状 [32, 7680])
    返回分段后的数据 (形状 [115, 32, 384])
    """
    n_segments = (trial_eeg.shape[-1] - window_samples) // step_samples + 1
    segments = np.zeros((n_segments, n_channels, window_samples))
    for seg_idx in range(n_segments):
        start = seg_idx * step_samples
        end = start + window_samples
        segments[seg_idx] = trial_eeg[:, start:end]
    return segments


def extract_all_trial(eeg_data, labels):
    # 初始化容器
    eeg_segments = np.zeros(
        (n_trials * n_segments_per_trial, n_channels, window_samples))
    label_segments = np.zeros((n_trials * n_segments_per_trial, n_labels))

    # 处理每个试次
    for trial_idx in range(n_trials):
        # 分割EEG
        trial_eeg = eeg_data[trial_idx]  # 形状 [32, 7680]
        segments = segment_trial(
            trial_eeg, window_samples, step_samples)  # [115, 32, 384]

        # 扩展标签
        trial_labels = labels[trial_idx]  # 形状 [4]
        expanded_labels = np.tile(
            trial_labels, (n_segments_per_trial, 1))  # [115, 4]

        # 填充到全局容器
        start_idx = trial_idx * n_segments_per_trial
        end_idx = start_idx + n_segments_per_trial
        eeg_segments[start_idx:end_idx] = segments
        label_segments[start_idx:end_idx] = expanded_labels

    return eeg_segments, label_segments


def calc_plv(eeg_data):
    analytic_signal = hilbert(eeg_data, axis=1)
    phases = np.angle(analytic_signal)

    # 计算PLV矩阵
    n_channels = phases.shape[0]
    plv_matrix = np.zeros((n_channels, n_channels))

    for i in range(n_channels):
        for j in range(n_channels):
            phase_diff = phases[i] - phases[j]
            plv = np.abs(np.mean(np.exp(1j * phase_diff)))
            plv_matrix[i, j] = plv

    return plv_matrix


if __name__ == "__main__":
    all_bands = {
        "delta": (0.1, 3),
        "theta": (4, 7),
        "low-alpha": (8, 9.5),
        "high-alpha": (10.5, 12),
        "alpha": (8, 12),
        "low_beta": (13, 16),
        "mid-beta": (17, 20),
        "high-beta": (21, 29),
        "beta": (13, 29),
        "gamma": (30, 50)
    }
    all_bands_dir = [os.path.join(
        'C:/Users/john/Desktop/learn-eeg/all_bands_data', band_name) for band_name in all_bands.keys()]

    for i in range(14, 33):
        eeg = []
        label = []
        for band_dir in all_bands_dir:
            eeg_data = np.load(
                f'{band_dir}/s{str(i).zfill(2)}.npy')
            labels = np.load(
                f'{band_dir}/s{str(i).zfill(2)}_labels.npy')

            eeg_segments, label_segments = extract_all_trial(eeg_data, labels)
            eeg.append(eeg_segments)
            label.append(label_segments)

        eeg = np.array(eeg)
        eeg = np.transpose(eeg, (1, 0, 2, 3))
        label = np.array(label[0])
        plv_mat_all_trial = []
        for trial_idx in range(eeg.shape[0]):
            eeg_data_10ch = eeg[trial_idx]
            plv_mat_trial = []
            for ch_idx in range(eeg.shape[1]):
                plv_mat_trial.append(calc_plv(eeg_data_10ch[ch_idx]))
            plv_mat_all_trial.append(plv_mat_trial)
            print(f'sub-----{i}trial----{trial_idx}')

        plv_mat_trial = np.array(plv_mat_all_trial)

        file_path = f'C:/Users/john/Desktop/learn-eeg/plv_matrix/s{str(i).zfill(2)}'
        os.makedirs(file_path, exist_ok=True)
        np.save(
            f'{file_path}/{i}.npy', plv_mat_all_trial)
        np.save(
            f'{file_path}/labels.npy', label)
